#!/usr/bin/python
# -*- coding: utf-8 -*-

# import sys
# import onnx
# from onnx import shape_inference
# input_onnx = sys.argv[1]
# output_onnx = sys.argv[2]
# model = onnx.load(input_onnx)
# model_with_shape = shape_inference.infer_shapes(model)
# onnx.save(model_with_shape, output_onnx)

import argparse
from auto_optimizer import OnnxGraph

def parse_input_shape(shape_str=None):
    if not shape_str:
        return {}
    inputs = shape_str.strip().split(';')
    shape_dict = {}
    for s in inputs:
        fields = s.split(':')
        assert len(fields) == 2
        dims = []
        for x in fields[1].split(','):
            try:
                x_ = int(x)
                if x_ > 0:
                    dims.append(x_)
                else:
                    dims.append(x)
            except:
                dims.append(x)
        shape_dict[fields[0]] = dims
    return shape_dict

def main():
    parser = argparse.ArgumentParser(
                        description='modify onnx model.')
    parser.add_argument('input_onnx', type=str, help='path to input onnx file.')
    parser.add_argument('output_onnx', type=str, 
                        help='path to save modified onnx model.')
    parser.add_argument('--input_shape', type=str, default=None,
                        help='eg: image:-1,3,224,224')
    args = parser.parse_args()

    shape_dict = parse_input_shape(args.input_shape)

    g = OnnxGraph.parse(args.input_onnx)
    input_index = {inp.name: i for i, inp in enumerate(g.inputs)}
    for input_name, dims in shape_dict.items():
        g.inputs[input_index[input_name]].shape = dims
    g.infershape()
    g.save(args.output_onnx)

    print('Done.')

if __name__ == "__main__":
    main()

